import numpy as np
import cv2

iou_threshold = 0.5
reward_movement_action = 1.0
reward_terminal_action = 3.0


# 注：img : [H->y, W->x, C]
class Action(object):
    def __init__(self, img, gts, classes_gt_objects, CLASS_OBJECT, scale_subregion=0.75):
        self.img = img  # 原始图片
        self.gts = gts  # 原始图片中的一个object
        self.classes_gt_objects = classes_gt_objects
        self.CLASS_OBJECT = CLASS_OBJECT
        # self.position 代表框的位置。下值为起始框的位置
        self.x = self.img.shape[1]
        self.y = self.img.shape[0]

        self.position = [0, 0, self.img.shape[1], self.img.shape[0]]
        self.scale_subregion = scale_subregion

    @property
    def now_img(self):  # 指 当前的 region_image，将其输入 feature_extractor 进而生成 state
        # print('original_image.shape = ', self.img.shape)
        # print('self.position = ', self.position)
        return self.img[self.position[1]:self.position[3], self.position[0]:self.position[2], :]

    @property
    def region_position(self):
        return self.position

    @property
    def actions(self):  # 可以遍历的方式调用各个 action
        return [self.top_left, self.top_right, self.bottom_left, self.bottom_right,
                self.middle, self.trigger]

    def calculate_iou(self, gt_mask, pos):
        gt_mask *= 1.0
        img_mask = np.zeros([self.y, self.x])
        img_mask[pos[1]:pos[3], pos[0]:pos[2]] = 1.0
        img_and = cv2.bitwise_and(img_mask, gt_mask)
        img_or = cv2.bitwise_or(img_mask, gt_mask)
        j = np.count_nonzero(img_and)
        i = np.count_nonzero(img_or)
        iou = float(float(j) / float(i))
        return iou

    def calculate_max_iou(self, pos=None):
        if not pos:
            pos = self.position
        max_iou = 0.0
        for gt_idx in range(self.gts.shape[2]):
            if self.classes_gt_objects[gt_idx] != self.CLASS_OBJECT:
                continue
            gt = self.gts[:, :, gt_idx]
            iou = self.calculate_iou(gt, pos)
            if iou > max_iou:
                max_iou = iou
        return max_iou

    # def best_action(self):
    #     before = self.score()  # 先求当前state下的IoU
    #     # 对 当前region 依次实施 9个action，并得到它们的 reward
    #     reward = [-before] * 9
    #     for i in range(9):
    #         score, pos = self.actions[i]()
    #         reward[i] = reward_movement_action if reward[i]+score > 0 else -1*reward_movement_action
    #         if i == 8:
    #             reward[i] = reward_terminal_action if score >= iou_threshold else -1*reward_terminal_action
    #     return np.array(reward)

    def action(self, num):  # 执行该函数，改变 region_image 的框
        iou, pos = self.actions[num]()
        self.position = pos
        return iou

    def top_left(self):
        pos = self.position[:]
        # if pos[3] - pos[1] > 18
        # if pos[2] - pos[0] > 20 and pos[3] - pos[1] > 18:
        region_x = int(self.scale_subregion * (pos[2] - pos[0]))
        region_y = int(self.scale_subregion * (pos[3] - pos[1]))
        pos[2] = pos[0] + region_x
        pos[3] = pos[1] + region_y
        return self.calculate_max_iou(pos), pos

    def top_right(self):
        pos = self.position[:]
        region_x = int(self.scale_subregion * (pos[2] - pos[0]))
        region_y = int(self.scale_subregion * (pos[3] - pos[1]))
        pos[0] = pos[2] - region_x
        pos[3] = pos[1] + region_y
        return self.calculate_max_iou(pos), pos

    def bottom_right(self):
        pos = self.position[:]
        region_x = int(self.scale_subregion * (pos[2] - pos[0]))
        region_y = int(self.scale_subregion * (pos[3] - pos[1]))
        pos[0] = pos[2] - region_x
        pos[1] = pos[3] - region_y
        return self.calculate_max_iou(pos), pos

    def bottom_left(self):
        pos = self.position[:]
        region_x = int(self.scale_subregion * (pos[2] - pos[0]))
        region_y = int(self.scale_subregion * (pos[3] - pos[1]))
        pos[2] = pos[0] + region_x
        pos[1] = pos[3] - region_y
        return self.calculate_max_iou(pos), pos

    def middle(self):
        pos = self.position[:]
        offset_x = int((pos[2] - pos[0]) * (1-self.scale_subregion) / 2)
        offset_y = int((pos[3] - pos[1]) * (1-self.scale_subregion) / 2)
        pos[0] = pos[0] + offset_x
        pos[1] = pos[1] + offset_y
        pos[2] = pos[2] - offset_x
        pos[3] = pos[3] - offset_y
        return self.calculate_max_iou(pos), pos

    def trigger(self):
        return self.calculate_max_iou(self.position), self.position
